#include <iostream>
#include <fstream>
#include <string>
#include <cstdio>  // for getchar
#include <cstring> // for memset

//===============================================

#include "kglobal.h"

#ifdef _FULLPROGRAM
#include "kcompile.h"
#endif

#include "kprogpath.h"
#include "kutils.h"
#include "kvm.h"

//===============================================

#define PROG_VERSION "v2.0"
#define SPEC_VERSION "11.01"
#define SPEC_VERSIONB "1101"

using namespace std;

#ifdef _FULLPROGRAM

#ifdef IS_WINDOWS
#define NATIVE_EXE "kimln.exe"
#define NATIVE_EXT ".exe"
#else
#define NATIVE_EXE "kimln"
#define NATIVE_EXT ".out"
#endif

bool compileOnly;  // -c/-x
enum _INPUT_FILE_T
{
	SrcFile, // source file *.kim
	CmpFile  // compiled file *.kc
} inputFile;
string fileName;
string fileNameWoExt;

string outFileName;
bool hasOutFile;

enum CommandOptions
{
	OPT_INVALID   = -1,
	OPT_KIM_RUN   = 0,  // run .kim
	OPT_KIM_CMP   = 1,  // compile to .kc
	OPT_KC_RUN    = 2,  // run .kc
	OPT_NOTHING   = 3,  // do nothing
	OPT_KIM_CMPEX = 4   // compile to native executable
};

void ShowHelp();

CommandOptions ParseArguments(int argc, char* argv[]);
string GetExt(const string &fn);
string GetFileNameWoExt(const string &fn);

bool CompileFile();
bool CompileAndRun();
bool CompileNative();
bool RunFile();
void Interact();

#endif //_FULLPROGRAM

bool RunCode(unsigned char *bin);
bool HasEmbeddedCode(unsigned char *&embedCode);

//===============================================

int main(int argc, char* argv[])
{
#ifdef _MSC_VER
	_CrtSetDbgFlag( _CRTDBG_ALLOC_MEM_DF | _CRTDBG_LEAK_CHECK_DF );
#endif

	GlobalInit();
	KVM::Init();

	unsigned char *bin;
	if (HasEmbeddedCode(bin))
	{
		if (!RunCode(bin))
		{
			delete [] bin;
			return 1;
		}

		delete [] bin;
	}

#ifdef _FULLPROGRAM
	else
	{
		switch (ParseArguments(argc, argv))
		{
		case OPT_INVALID:
#ifdef _DEBUG
			getchar();
#endif
			return -1;
		case OPT_KIM_RUN:
			// compile and run
			CompileAndRun();
			break;
		case OPT_KIM_CMP:
			// compile only
			if (CompileFile())
				cout << "Compiled successfully to " << outFileName << endl;
			else
				cout << "Compiled unsuccessfully." << endl;
			break;
		case OPT_KC_RUN:
			// run compiled file
			RunFile();
			break;
		case OPT_NOTHING:
			// nothing given, switch to interactive mode
			Interact();
			break;
		case OPT_KIM_CMPEX:
			// compile to native executable
			if (CompileNative())
				cout << "Compiled successfully to " << outFileName << endl;
			else
				cout << "Compiled unsuccessfully." << endl;
			break;
		}

	} //HasEmbeddedCode?

#endif //_FULLPROGRAM

	KVM::CleanUp();
	GlobalCleanUp();

#if _DEBUG
	getchar();
#endif
	
#ifdef _MSC_VER
	//_CrtDumpMemoryLeaks();
#endif

	return 0;
}

#ifdef _FULLPROGRAM

void ShowHelp()
{
	cout << "KimL compiler and interpreter - " << PROG_VERSION << endl;
	cout << "KimL Specification " << SPEC_VERSION << endl;
	cout << "Usage: kiml [-c] [-d] [-h] [-o output] file" << endl << endl;

	cout << "Options:" << endl;
	cout << "\t-c\tCompile but not run.\n\t\tOnly applicable for *.kim file." << endl;
	cout << "\t-x\tCompile to native executable file.\n\t\tOnly applicable for *.kim file." << endl;
	cout << "\t-d\tEmbed debug information in the compiled file.\n\t\tOnly applicable for *.kim file." << endl;
	cout << "\t-h\tShow help message." << endl;
	cout << "\t-o\tSpecify output file." << endl << endl;
}

CommandOptions ParseArguments(int argc, char* argv[])
{
	bool hasFile = false;
	bool hasHelp = false;
	bool compileNative = false;

	for (int i = 1; i < argc; ++i)
	{
		if (my_strequ(argv[i], "-c"))
		{
			if (!compileOnly)
			{
				compileOnly = true;
			}
			else
			{
				cout << "*** Duplicate option -c/-x." << endl;
				return OPT_INVALID;
			}
		}
		else if (my_strequ(argv[i], "-x"))
		{
			if (!compileOnly)
			{
				compileOnly = true;
				compileNative = true;
			}
			else
			{
				cout << "*** Duplicate option -c/-x." << endl;
				return OPT_INVALID;
			}
		}
		else if (my_strequ(argv[i], "-d"))
		{
			if (!hasDebugInfo)
			{
				hasDebugInfo = true;
			}
			else
			{
				cout << "*** Duplicate option -d." << endl;
				return OPT_INVALID;
			}
		}
		else if (my_strequ(argv[i], "-h"))
		{
			if (!hasHelp)
			{
				hasHelp = true;
			}
			else
			{
				cout << "*** Duplicate option -h." << endl;
				return OPT_INVALID;
			}
		}
		else if (my_strequ(argv[i], "-o"))
		{
			if (!hasOutFile)
			{
				hasOutFile = true;

				if (i == argc - 1)
				{
					cout << "*** Output filename not specified." << endl;
					return OPT_INVALID;
				}
				else
				{
					outFileName = argv[++i];
				}
			}
			else
			{
				cout << "*** Duplicate output file." << endl;
				return OPT_INVALID;
			}
		}
		else
		{
			if (!hasFile)
			{
				hasFile = true;
				fileName = argv[i];
				
				string ext = GetExt(fileName);
				
				if (ext.size() == 0)
				{
					cout << "*** Unknown file type." << endl;
					return OPT_INVALID;
				}
				
				if (my_strequ(ext.c_str(), "kim"))
				{
					inputFile = SrcFile;
					fileNameWoExt = GetFileNameWoExt(fileName);
				}
				else if (my_strequ(ext.c_str(), "kc"))
				{
					inputFile = CmpFile;
				}
				else
				{
					cout << "*** Unknown file type: " << ext << endl;
					return OPT_INVALID;
				}
			}
			else
			{
				cout << "*** Duplicate input file." << endl;
				return OPT_INVALID;
			}
		}
	}

	if (hasHelp)
		ShowHelp();

	if (!hasFile)
	{
		if (!compileOnly && !hasDebugInfo)
		{
			return OPT_NOTHING;
		}
		else
		{
			cout << "*** Options given without a file." << endl;
			return OPT_INVALID;
		}
	}

	if (compileOnly)
	{
		if (inputFile == SrcFile)
		{
			return compileNative ? OPT_KIM_CMPEX : OPT_KIM_CMP;
		}
		else
		{
			cout << "*** Input file must be a *.kim file." << endl;
			return OPT_INVALID;
		}
	}
	else
	{
		if (inputFile == SrcFile)
			return OPT_KIM_RUN;
		else
			return OPT_KC_RUN;
	}
}

string GetExt(const string &fn)
{
	int len = fn.length();
	for (int i = len - 1; i >= 0; --i)
		if (fn[i] == '.' && i < len-1)
			return fn.substr(i+1);

	return string("");
}

string GetFileNameWoExt(const string &fn)
{
	int len = fn.length();
	for (int i = len - 1; i >= 0; --i)
		if (fn[i] == '.')
			return fn.substr(0, i);

	return string("");
}

bool CompileFile()
{
	static char magic[4] = { 'k', 'i', 'm', 'l' };
	static char ver[4] =   { '1', '1', '0', '1' };
	static char pad[8] = { 0 };

	ifstream ifs(fileName.c_str(), ios::binary);
	
	if (!ifs.is_open())
	{
		cout << "*** Cannot open file: " << fileName << endl;
		return false;
	}

	k_uint len;
	ifs.seekg(0, ios::end);
	len = (k_uint)ifs.tellg();
	ifs.seekg(0, ios::beg);

	char *src = new char[len + 1];
	memset(src, 0, len + 1);
	ifs.read(src, len);
	ifs.close();

	if (!Compile(src))
	{
		delete[] src;
		return false;
	}

	delete[] src;

	if (!hasOutFile)
		outFileName = fileNameWoExt.append(".kc");

	ofstream ofs(outFileName.c_str(), ios::binary);

	if (!ofs.is_open())
	{
		cout << "*** Cannot write binary file: " << fileNameWoExt << endl;
		bytecode->CleanUp();
		strPool->CleanUp();
		return false;
	}

	// Header
	ofs.write(magic, sizeof(magic));
	ofs.write(ver, sizeof(ver));
	ofs.write(pad, sizeof(pad));

	// Bytecode section
	const unsigned char *bc = bytecode->GetStream();
	ofs.write((const char *)&programSize, sizeof(k_int));
	ofs.write((const char *)bc, programSize);

	// String pool section
	const char *sp = strPool->GetPool();
	ofs.write((const char *)&strPoolSize, sizeof(k_int));
	ofs.write(sp, strPoolSize);

	ofs.close();

	bytecode->CleanUp();
	strPool->CleanUp();

	return true;
}

bool CompileAndRun()
{
	ifstream ifs(fileName.c_str(), ios::binary);
	
	if (!ifs.is_open())
	{
		cout << "*** Cannot open file: " << fileName << endl;
		return false;
	}

	k_uint len;
	ifs.seekg(0, ios::end);
	len = (k_uint)ifs.tellg();
	ifs.seekg(0, ios::beg);

	char *src = new char[len + 1];
	memset(src, 0, len + 1);
	ifs.read(src, len);
	ifs.close();

	if (!Compile(src))
	{
		delete[] src;
		bytecode->CleanUp();
		strPool->CleanUp();
		return false;
	}

	delete[] src;

	bool result = KVM::Run(bytecode->GetStream(), programSize, strPool->GetPool());

	bytecode->CleanUp();
	strPool->CleanUp();

	return result;
}

bool CompileNative()
{
	static char magic[4] = { 'k', 'i', 'm', 'l' };
	static char ver[4] =   { '1', '1', '0', '1' };
	static char pad[8] = { 0 };
	static unsigned char embedMagic[4] = { 0x0A, 0xCE, 0xCA, 0xDE };

	k_uint offset;

	ifstream ifs(fileName.c_str(), ios::binary);
	
	if (!ifs.is_open())
	{
		cout << "*** Cannot open file: " << fileName << endl;
		return false;
	}

	k_uint len;
	ifs.seekg(0, ios::end);
	len = (k_uint)ifs.tellg();
	ifs.seekg(0, ios::beg);

	char *src = new char[len + 1];
	memset(src, 0, len + 1);
	ifs.read(src, len);
	ifs.close();

	if (!Compile(src))
	{
		delete[] src;
		return false;
	}

	delete[] src;
	
	ifs.open(NATIVE_EXE, ios::binary);
	if (!ifs.is_open())
	{
		cout << "*** " << NATIVE_EXE << " not found. Use this program." << endl;

		char *progPath = GetProgramPath();
		ifs.open(progPath, ios::binary);

		delete [] progPath;

		if (!ifs.is_open())
		{
			cout << "*** Cannot open this program." << endl;
			bytecode->CleanUp();
			strPool->CleanUp();
			return false;
		}
	}

	ifs.seekg(0, ios::end);
	offset = (k_uint)ifs.tellg();
	ifs.seekg(0, ios::beg);

	if (!hasOutFile)
		outFileName = fileNameWoExt.append(NATIVE_EXT);

	ofstream ofs(outFileName.c_str(), ios::binary);

	if (!ofs.is_open())
	{
		cout << "*** Cannot write binary file: " << fileNameWoExt << endl;
		ifs.close();
		bytecode->CleanUp();
		strPool->CleanUp();
		return false;
	}

	// bootstrap
	ofs << ifs.rdbuf();
	ifs.close();

	// Header
	ofs.write(magic, sizeof(magic));
	ofs.write(ver, sizeof(ver));
	ofs.write(pad, sizeof(pad));

	// Bytecode section
	const unsigned char *bc = bytecode->GetStream();
	ofs.write((const char *)&programSize, sizeof(k_int));
	ofs.write((const char *)bc, programSize);

	// String pool section
	const char *sp = strPool->GetPool();
	ofs.write((const char *)&strPoolSize, sizeof(k_int));
	ofs.write(sp, strPoolSize);

	// Embed information
	ofs.write((const char *)&offset, sizeof(k_uint));
	ofs.write((const char *)embedMagic, sizeof(embedMagic));

	ofs.close();

	bytecode->CleanUp();
	strPool->CleanUp();

	return true;
}

bool RunFile()
{
	ifstream ifs(fileName.c_str(), ios::binary);
	
	ifs.seekg(0, ios::end);
	k_uint binSize = (k_uint)ifs.tellg();

	unsigned char *bin = new unsigned char[binSize];
	ifs.seekg(0, ios::beg);
	ifs.read((char *)bin, binSize);

	bool result = RunCode(bin);
	
	delete [] bin;

	return result;
}

void Interact()
{
	k_string src;
	char ln[1024]; //line buffer
	int c, n;
	k_uint lno = 1;

	cout << "KimL Interactive Mode.\n"
		 << "Enter lines of code and terminate using\n\
	Ctrl+Z on a new line (Windows) or\n\
	Ctrl+D (Linux).\n"
		 << "Type '.exit' to close the program.\n\n";

	for (;;)
	{
		for (;;)
		{
			printf("%03u> ", lno);

			for (n = 0; n < 1023; ++n)
			{
				c = fgetc(stdin);
				if (c == '\n')
				{
					ln[n] = c;
					ln[n+1] = 0;
					break;
				}
				else if (c == EOF)
				{
					break;
				}

				ln[n] = c;
			}
			
			++lno;

			if (my_strequ(ln, ".exit\n"))
			{
				src.clear();
				return;
			}

			if (c == EOF /*Ctrl^Z*/)
				break;

			src.append(ln);
		}

		if (Compile(src.c_str()))
			KVM::Run(bytecode->GetStream(), programSize, strPool->GetPool());

		bytecode->CleanUp();
		strPool->CleanUp();

		src.clear();
		lno = 1;
	}
}

#endif //_FULLPROGRAM

bool RunCode(unsigned char *bin)
{
	const unsigned char *code;
	const char *pool;
	k_uint codeSize;
	k_uint poolSize;
	k_uint pos = 0;

	char magic[5] = { 0 };
	char ver[5] = { 0 };
	char pad[8] = { 0 };

	memcpy(magic, bin + pos, 4);
	pos += 4;
	if (!my_strequ(magic, "kiml"))
	{
		cout << "Invalid binary file." << endl;
#ifdef _DEBUG
		cout << "Wrong magic." << endl;
#endif
		return false;
	}

	memcpy(ver, bin + pos, 4);
	pos += 4;
	if (!my_strequ(ver, SPEC_VERSIONB))
	{
		cout << "Invalid binary file." << endl;
#ifdef _DEBUG
		cout << "Wrong version." << endl;
#endif
		return false;
	}

	memcpy(pad, bin + pos, 8);
	pos += 8;
	for (int i = 0; i < 8; ++i)
	{
		if (pad[i])
		{
			cout << "Invalid binary file." << endl;
#ifdef _DEBUG
			cout << "Wrong header padding." << endl;
#endif
			return false;
		}
	}

	memcpy(&codeSize, bin + pos, sizeof(k_uint));
	pos += sizeof(k_uint);

	code = new unsigned char[codeSize];
	memcpy((unsigned char *)code, bin + pos, codeSize);
	pos += codeSize;

	memcpy(&poolSize, bin + pos, sizeof(k_uint));
	pos += sizeof(k_uint);

	pool = new char[poolSize];
	memcpy((char *)pool, bin + pos, poolSize);

	bool result = KVM::Run(code, codeSize, pool);

	delete [] code;
	delete [] pool;

	return result;
}

bool HasEmbeddedCode(unsigned char *&embedCode)
{
	char *progPath = GetProgramPath();
	ifstream ifs(progPath, ios::binary);

	if (!ifs.is_open())
	{
		cout << "This program cannot read itself. Path: " << progPath << endl;
		delete []progPath;
		return false;
	}

	delete [] progPath;

	ifs.seekg(-4, ios::end);

	char embedMagic[4] = { 0 };
	ifs.read(embedMagic, 4);

	if (embedMagic[0] != '\x0A'
		|| embedMagic[1] != '\xCE'
		|| embedMagic[2] != '\xCA'
		|| embedMagic[3] != '\xDE')
	{
		return false;
	}

	ifs.seekg(0, ios::end);
	k_uint fileSize = (k_uint)ifs.tellg();

	ifs.seekg(-4 - (signed)sizeof(k_uint), ios::end);

	k_uint offset;
	ifs.read((char *)&offset, sizeof(k_uint));
	
	k_uint embedSize = fileSize - offset - sizeof(embedMagic) - sizeof(k_uint);
	if (embedSize <= 0)
	{
		cout << "Program corrupted." << endl;
		return false;
	}

	embedCode = new unsigned char[embedSize];
	
	ifs.seekg(offset, ios::beg);
	ifs.read((char *)embedCode, embedSize);

	ifs.close();

	return true;
}
